package com.ww.exercise.channel.channels

import kotlinx.coroutines.*
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.selects.select
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import java.util.*
import kotlin.coroutines.CoroutineContext
import kotlin.time.Duration
import kotlin.time.Duration.Companion.seconds

/**
 * 设计一个基于 Channel 的任务调度系统，支持不同优先级的后台任务处理。要求：
 *
 * 高优先级任务应优先于低优先级任务执行
 * 实现任务队列的背压控制，当队列满时能够合理拒绝新任务
 * 支持任务的暂停和恢复功能
 * 能够优雅地关闭调度器，确保所有已提交任务完成
 */

// 任务优先级枚举
enum class TaskPriority {
    HIGH,
    MEDIUM,
    LOW
}

// 任务结果密封类
sealed class TaskResult {
    data class Success(val taskId: String) : TaskResult()
    data class Failure(val taskId: String, val exception: Exception) : TaskResult()
    data class Rejected(val taskId: String, val reason: String) : TaskResult()
}

// 任务数据类
data class Task(
    val id: String = UUID.randomUUID().toString(),
    val priority: TaskPriority,
    val delay: Duration = Duration.ZERO,
    val action: suspend () -> Unit
)

/**
 * 基于Channel的优先级任务调度器
 * 支持任务优先级、背压控制、暂停/恢复和优雅关闭
 */
class PriorityTaskScheduler(
    private val highPriorityCapacity: Int = 100,
    private val mediumPriorityCapacity: Int = 200,
    private val lowPriorityCapacity: Int = 500,
    private val maxWorkers: Int = 4,
    private val coroutineContext: CoroutineContext = Dispatchers.Default
) {
    // 不同优先级的任务通道
    private val highPriorityChannel = Channel<Task>(highPriorityCapacity)
    private val mediumPriorityChannel = Channel<Task>(mediumPriorityCapacity)
    private val lowPriorityChannel = Channel<Task>(lowPriorityCapacity)

    // 结果通道
    private val resultChannel = Channel<TaskResult>(Channel.UNLIMITED)

    // 控制状态的变量
    private val stateMutex = Mutex()
    private var isPaused = false
    private var isShutdown = false
    private var workers: List<Job> = emptyList()
    private val scope = CoroutineScope(coroutineContext + SupervisorJob())

    // 启动调度器
    fun start() {
        scope.launch {
            stateMutex.withLock {
                if (workers.isNotEmpty()) return@launch

                workers = (1..maxWorkers).map { workerId ->
                    launch {
                        workerLoop(workerId)
                    }
                }
            }
        }
    }

    /**
     * 提交任务到调度器
     */
    suspend fun submitTask(task: Task): TaskResult {
        return withContext(coroutineContext) {
            stateMutex.withLock {
                if (isShutdown) {
                    return@withContext TaskResult.Rejected(task.id, "调度器已关闭，无法接受新任务")
                }

                // 根据优先级发送到不同的通道，处理背压
                val sendResult = when (task.priority) {
                    TaskPriority.HIGH -> highPriorityChannel.trySend(task)
                    TaskPriority.MEDIUM -> mediumPriorityChannel.trySend(task)
                    TaskPriority.LOW -> lowPriorityChannel.trySend(task)
                }

                if (sendResult.isSuccess) {
                    TaskResult.Success(task.id)
                } else {
                    TaskResult.Rejected(task.id, "${task.priority}优先级任务队列已满")
                }
            }
        }
    }

    /**
     * 暂停任务调度
     */
    suspend fun pause() {
        stateMutex.withLock {
            isPaused = true
        }
    }

    /**
     * 恢复任务调度
     */
    suspend fun resume() {
        stateMutex.withLock {
            isPaused = false
        }
    }

    /**
     * 优雅关闭调度器，等待所有任务完成
     */
    suspend fun shutdown() {
        stateMutex.withLock {
            if (isShutdown) return@withLock

            isShutdown = true
            // 关闭所有输入通道，不再接受新任务
            highPriorityChannel.close()
            mediumPriorityChannel.close()
            lowPriorityChannel.close()

            // 等待所有工作协程完成
            workers.joinAll()

            // 取消作用域
            scope.cancel()
        }
    }

    /**
     * 订阅任务结果
     */
    fun subscribeToResults(handler: (TaskResult) -> Unit): Job {
        return scope.launch {
            for (result in resultChannel) {
                handler(result)
            }
        }
    }

    /**
     * 工作协程循环，处理任务
     */
    @OptIn(DelicateCoroutinesApi::class)
    private suspend fun workerLoop(workerId: Int) {
        try {
            while (true) {
                // 检查是否需要暂停
                checkPauseState()

                // 使用select表达式实现优先级选择
                val task = select<Task?> {
                    // 高优先级任务优先处理
                    highPriorityChannel.onReceiveCatching { it.getOrNull() }
                    // 其次是中优先级
                    mediumPriorityChannel.onReceiveCatching { it.getOrNull() }
                    // 最后是低优先级
                    lowPriorityChannel.onReceiveCatching { it.getOrNull() }
                }

                // 如果所有通道都关闭且没有任务了，退出循环
                if (task == null) {
                    if (highPriorityChannel.isClosedForReceive &&
                        mediumPriorityChannel.isClosedForReceive &&
                        lowPriorityChannel.isClosedForReceive
                    ) {
                        break
                    }
                    continue
                }

                // 执行任务
                try {
                    // 处理任务延迟
                    if (task.delay > Duration.ZERO) {
                        delay(task.delay)
                    }

                    // 再次检查是否在延迟后被暂停
                    checkPauseState()

                    // 执行任务
                    task.action()
                    resultChannel.send(TaskResult.Success(task.id))
                } catch (e: Exception) {
                    resultChannel.send(TaskResult.Failure(task.id, e))
                }
            }
        } catch (e: CancellationException) {
            // 正常取消，忽略
        } catch (e: Exception) {
            // 记录工作协程异常
            println("Worker $workerId 异常: ${e.message}")
        }
    }

    /**
     * 检查暂停状态，如果已暂停则挂起直到恢复
     */
    private suspend fun checkPauseState() {
        while (true) {
            stateMutex.withLock {
                if (!isPaused) return@checkPauseState
            }
            // 短暂延迟后再次检查，避免频繁加锁
            delay(100)
        }
    }

    /**
     * 获取当前队列状态
     */
    suspend fun getQueueStatus(): Map<TaskPriority, Int> {
        return stateMutex.withLock {
            mapOf(
                TaskPriority.HIGH to countChannelElements(highPriorityChannel),
                TaskPriority.MEDIUM to countChannelElements(mediumPriorityChannel),
                TaskPriority.LOW to countChannelElements(lowPriorityChannel)
            )
        }
    }

    /**
     * 工具函数：安全统计缓冲Channel中的元素数量（非阻塞、无数据丢失）
     * 原理：暂存所有元素→计数→放回元素，全程加锁保证线程安全
     */
    private suspend fun countChannelElements(channel: Channel<Task>): Int {
        val tempList = mutableListOf<Task>()
        // 非阻塞循环取出所有元素（tryReceive()无元素时返回Failure）
        while (true) {
            val receiveResult = channel.tryReceive()
            if (receiveResult.isFailure) break // 无更多元素，退出循环
            tempList.add(receiveResult.getOrThrow()) // 暂存元素
        }
        // 将暂存的元素放回Channel（保证顺序不变）
        tempList.forEach { channel.send(it) }
        // 返回统计数量
        return tempList.size
    }
}

fun main() = runBlocking {
    val scheduler = PriorityTaskScheduler(
        highPriorityCapacity = 5,
        mediumPriorityCapacity = 10,
        lowPriorityCapacity = 20,
        maxWorkers = 2
    )

    scheduler.subscribeToResults { result ->
        when (result) {
            is TaskResult.Success -> println("任务 ${result.taskId.substring(0, 8)} 执行成功")
            is TaskResult.Failure -> println("任务 ${result.taskId.substring(0, 8)} 执行失败: ${result.exception.message}")
            is TaskResult.Rejected -> println("任务 ${result.taskId.substring(0, 8)} 被拒绝: ${result.reason}")
        }
    }

    scheduler.start()
    println("调度器启动")

    // 【重点修复】提交测试任务：提前定义 id 和 priority，避免自引用
    repeat(10) { i ->
        // 1. 提前生成任务ID（与 Task 默认的 UUID 逻辑一致）
        val taskId = UUID.randomUUID().toString()
        // 2. 提前确定任务优先级（逻辑不变）
        val taskPriority = when {
            i % 5 == 0 -> TaskPriority.HIGH
            i % 3 == 0 -> TaskPriority.MEDIUM
            else -> TaskPriority.LOW
        }
        // 3. 创建 Task 时，使用提前定义的 taskId 和 taskPriority，而非引用未初始化的 task
        val task = Task(
            id = taskId, // 显式传入提前生成的ID
            priority = taskPriority, // 显式传入提前确定的优先级
            delay = if (i % 4 == 0) 1.seconds else Duration.ZERO
        ) {
            // 直接使用提前定义的 taskId 和 taskPriority，避免引用 task
            println("执行任务 ${taskId.substring(0, 8)} (优先级: $taskPriority)")
            delay(500) // 模拟任务执行时间
        }

        val result = scheduler.submitTask(task)
        if (result is TaskResult.Rejected) {
            println("提交任务 ${taskId.substring(0, 8)} 失败: ${result.reason}")
        }
    }

    // 【同步修复】暂停期间提交任务的逻辑（同样避免自引用）
    delay(2000)
    println("暂停调度器")
    scheduler.pause()

    // 修复暂停任务的创建：提前定义 id 和 priority
    val pausedTaskId = UUID.randomUUID().toString()
    val pausedTask = Task(
        id = pausedTaskId,
        priority = TaskPriority.MEDIUM
    ) {
        // 使用提前定义的 pausedTaskId，而非引用 pausedTask
        println("这个任务在暂停期间提交，应该在恢复后执行: ${pausedTaskId.substring(0, 8)}")
    }
    scheduler.submitTask(pausedTask)

    // 恢复调度器
    delay(2000)
    println("恢复调度器")
    scheduler.resume()

    // 等待一段时间后关闭调度器
    delay(3000)
    println("开始关闭调度器")
    scheduler.shutdown()
    println("调度器已关闭")
}